"""
手写笔记识别整理系统 - Streamlit Web版
基于文心多模态大模型 ERNIE-4.5-VL-28B-A3B
上传图片 -> AI 识别 -> 返回整理后的笔记
"""

import os
import json
import base64
import io
import cv2
import numpy as np
from PIL import Image, ImageEnhance
from openai import OpenAI
import streamlit as st
import time
import hashlib
from typing import Dict, Any, Optional, List
import logging

# PDF处理相关库
try:
    import fitz  # PyMuPDF
    PDF_AVAILABLE = True
except ImportError:
    PDF_AVAILABLE = False
    st.warning("PDF功能需要安装PyMuPDF: pip install PyMuPDF")

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ------------------ 内核：NoteRecognition ------------------
class NoteRecognition:
    def __init__(self):
        self.client = OpenAI(
            api_key="7b97b5e65d1248169aab2d56f67d2b0fbcb146a2",
            base_url="https://aistudio.baidu.com/llm/lmapi/v3",
        )
        self.ocr = None
        self.retry_count = 3
        self.session_cache = {}
        self._init_ocr()

    def _init_ocr(self):
        """初始化OCR引擎"""
        try:
            from paddleocr import PaddleOCR
            self.ocr = PaddleOCR(
                use_textline_orientation=False, 
                lang='ch',
                # show_log=False,  # 减少日志输出
                # use_gpu=False    # 使用CPU模式，更稳定
            )
            logger.info("OCR引擎初始化成功")
        except Exception as e:
            logger.warning(f"OCR引擎初始化失败: {e}")
            self.ocr = None

    def process_pdf(self, pdf_file) -> List[Image.Image]:
        """处理PDF文件，返回页面图像列表"""
        if not PDF_AVAILABLE:
            st.error("PDF功能不可用，请安装PyMuPDF: pip install PyMuPDF")
            return []
        
        try:
            # 读取PDF文件
            pdf_document = fitz.open(stream=pdf_file.read(), filetype="pdf")
            images = []
            
            # 处理每一页
            for page_num in range(pdf_document.page_count):
                page = pdf_document.load_page(page_num)
                
                # 设置缩放因子以提高图像质量
                mat = fitz.Matrix(2.0, 2.0)  # 2倍缩放
                pix = page.get_pixmap(matrix=mat)
                
                # 转换为PIL图像
                img_data = pix.tobytes("png")
                img = Image.open(io.BytesIO(img_data))
                images.append(img)
                
                logger.info(f"PDF第{page_num + 1}页处理完成")
            
            pdf_document.close()
            return images
            
        except Exception as e:
            logger.error(f"PDF处理失败: {e}")
            st.error(f"PDF处理失败: {str(e)}")
            return []

    def check_image_quality(self, image_pil: Image.Image) -> Dict[str, Any]:
        """检查图片质量并给出建议"""
        quality_info = {
            "width": image_pil.size[0],
            "height": image_pil.size[1],
            "format": image_pil.format,
            "mode": image_pil.mode,
            "issues": [],
            "suggestions": []
        }
        
        # 检查尺寸
        if image_pil.size[0] < 400 or image_pil.size[1] < 400:
            quality_info["issues"].append("图片尺寸较小")
            quality_info["suggestions"].append("建议使用分辨率更高的图片")
        
        # 检查清晰度
        image_array = np.array(image_pil.convert('L'))
        laplacian_var = cv2.Laplacian(image_array, cv2.CV_64F).var()
        
        if laplacian_var < 50:
            quality_info["issues"].append("图片清晰度较低")
            quality_info["suggestions"].append("建议使用更清晰的图片")
        elif laplacian_var > 500:
            quality_info["issues"].append("图片可能过于锐化")
            quality_info["suggestions"].append("建议使用正常清晰度的图片")
        
        return quality_info

    def enhance_image(self, image_pil: Image.Image) -> Image.Image:
        """图像增强处理"""
        try:
            # 转换为RGB模式
            if image_pil.mode != 'RGB':
                image_pil = image_pil.convert('RGB')
            
            # 对比度增强
            enhancer = ImageEnhance.Contrast(image_pil)
            image_pil = enhancer.enhance(1.2)
            
            # 锐度增强
            enhancer = ImageEnhance.Sharpness(image_pil)
            image_pil = enhancer.enhance(1.1)
            
            logger.info("图像增强处理完成")
            return image_pil
        except Exception as e:
            logger.warning(f"图像增强失败: {e}")
            return image_pil

    def preprocess_image(self, image_pil: Image.Image) -> np.ndarray:
        """图像预处理"""
        try:
            # 图像增强
            enhanced_image = self.enhance_image(image_pil)
            
            # 转换为numpy数组
            image = np.array(enhanced_image)
            
            # 降噪处理（参数优化）
            image = cv2.fastNlMeansDenoisingColored(
                image, 
                None, 
                h=8,      # 降低强度，减少失真
                hColor=8, 
                templateWindowSize=7, 
                searchWindowSize=21
            )
            
            logger.info("图像预处理完成")
            return image
        except Exception as e:
            logger.error(f"图像预处理失败: {e}")
            return np.array(image_pil)

    def extract_text_with_ocr(self, image: np.ndarray) -> str:
        """OCR文字识别"""
        if not self.ocr:
            return ""
        
        try:
            # 转换为BGR格式
            image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            
            # OCR识别
            result = self.ocr.ocr(image_bgr, cls=False)
            
            if result and result[0]:
                texts = []
                for line in result[0]:
                    if len(line) >= 2 and len(line[1]) >= 2:
                        text = line[1][0]
                        confidence = line[1][1]
                        # 只保留置信度较高的文字
                        if confidence > 0.6:
                            texts.append(text)
                
                extracted_text = "\n".join(texts)
                logger.info(f"OCR识别成功，识别到 {len(texts)} 行文字")
                return extracted_text
            else:
                logger.warning("OCR未识别到任何文字")
                return ""
                
        except Exception as e:
            logger.error(f"OCR识别失败: {e}")
            return ""

    def image_to_base64(self, image: np.ndarray) -> str:
        """图像转base64编码"""
        try:
            pil_image = Image.fromarray(image)
            buffer = io.BytesIO()
            pil_image.save(buffer, format='PNG', optimize=True)
            return base64.b64encode(buffer.getvalue()).decode()
        except Exception as e:
            logger.error(f"图像转base64失败: {e}")
            return ""

    def _get_image_hash(self, image: np.ndarray) -> str:
        """生成图像哈希值"""
        try:
            image_bytes = self.image_to_base64(image).encode()
            return hashlib.md5(image_bytes).hexdigest()
        except Exception as e:
            logger.error(f"生成图像哈希失败: {e}")
            return str(time.time())  # 降级方案

    def analyze_note(self, image: np.ndarray, extracted_text: str = "") -> Dict[str, Any]:
        """AI分析手写笔记"""
        image_hash = self._get_image_hash(image)
        
        # 检查缓存
        if image_hash in self.session_cache:
            logger.info("使用缓存结果")
            return self.session_cache[image_hash]
        
        image_base64 = self.image_to_base64(image)
        if not image_base64:
            return {"错误": "图像编码失败"}
        
        # 优化的提示词 - 针对笔记识别
        prompt = """你是一位专业的高中语文教师，负责对作文进行评分和修改建议。请严格按照以下标准进行评分：
    一、基础等级（40分）
    1. 内容项（20分）
    A. 题意理解（5分）
        - 是否准确理解题目要求
        - 是否切中题目核心
        - 是否把握题目重点
    
    B. 中心思想（5分）
        - 是否明确表达中心
        - 是否贯穿全文
        - 是否深刻有见地
    
    C. 内容充实（5分）
        - 是否材料丰富
        - 是否论据充分
        - 是否事例典型
    
    D. 感情真实（5分）
        - 是否感情真挚
        - 是否表达自然
        - 是否打动人心

    2. 表达项（20分）
    A. 语言表达（7分）
        - 是否用词准确
        - 是否句式灵活
        - 是否语言流畅
        - 是否表达生动
    
    B. 结构完整（7分）
        - 是否层次分明
        - 是否过渡自然
        - 是否首尾呼应
        - 是否结构严谨
    
    C. 文体规范（6分）
        - 是否符合文体要求
        - 是否格式正确
        - 是否标点准确

    二、发展等级（20分）
    1. 深刻性（5分）
    - 是否透过现象看本质
    - 是否揭示事物内在联系
    - 是否观点具有启发性
    - 是否见解独到

    2. 丰富性（5分）
    - 是否材料丰富多样
    - 是否论据充分有力
    - 是否形象丰满生动
    - 是否意境深远

    3. 文采（5分）
    - 是否用词贴切优美
    - 是否句式灵活多变
    - 是否善用修辞手法
    - 是否文句富有表现力

    4. 创意（5分）
    - 是否见解新颖独特
    - 是否材料新鲜有趣
    - 是否构思精巧别致
    - 是否想象丰富合理

    三、其他评分项
    1. 错别字：每字扣1分，最多扣5分
    2. 标点符号：3处以上错误酌情扣分
    3. 字数不足：每少50字扣1分
    4. 无标题：扣2分

    请按照以下格式给出评分和建议：

    【总分：XX分】
    - 基础等级：XX分（内容：XX分，表达：XX分）
    - 发展等级：XX分
    - 其他扣分：XX分

    【具体评分理由】
    1. 内容方面：
    - 优点：
        * （具体列举2-3点）
    - 不足：
        * （具体列举2-3点）

    2. 表达方面：
    - 优点：
        * （具体列举2-3点）
    - 不足：
        * （具体列举2-3点）

    3. 发展等级：
    - 亮点：
        * （具体列举2-3点）
    - 提升空间：
        * （具体列举2-3点）

    【修改建议】
    1. 内容提升建议：
    - （具体建议2-3点）
    2. 表达优化建议：
    - （具体建议2-3点）
    3. 发展等级提升建议：
    - （具体建议2-3点）

    注意：
    1. 评分要客观公正，严格按照标准执行
    2. 建议要具体可行，避免空泛
    3. 语言要专业规范，体现教师专业性
    4. 要注重鼓励性评价，保护学生写作积极性"""

        if extracted_text:
            prompt += f"\n\nOCR识别到的文字内容：\n{extracted_text}"

        try:
            response = self.client.chat.completions.create(
                model="ernie-4.5-vl-28b-a3b",
                messages=[{
                    "role": "user", 
                    "content": [
                        {"type": "text", "text": prompt},
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}
                    ]
                }],
                stream=False,
                extra_body={"penalty_score": 1, "enable_thinking": True},
                max_completion_tokens=2000,
                temperature=0.1,  # 降低随机性
                top_p=0.9,
            )

            result_text = response.choices[0].message.content
            result = self._parse_ai_response(result_text)
            
            # 缓存成功的结果
            if "错误" not in result:
                self.session_cache[image_hash] = result
                logger.info("结果已缓存")
            
            return result
                
        except Exception as e:
            logger.error(f"API调用失败: {e}")
            return {"错误": f"API调用失败: {str(e)}"}

    def _parse_ai_response(self, response_text: str) -> Dict[str, Any]:
        """智能解析AI响应"""
        # 尝试直接解析
        try:
            return json.loads(response_text)
        except json.JSONDecodeError:
            pass
        
        # 尝试提取JSON部分
        try:
            start_idx = response_text.find('{')
            end_idx = response_text.rfind('}')
            
            if start_idx != -1 and end_idx != -1 and end_idx > start_idx:
                json_part = response_text[start_idx:end_idx + 1]
                # 清理特殊字符
                json_part = json_part.replace('\n', ' ').replace('\r', ' ')
                return json.loads(json_part)
        except json.JSONDecodeError:
            pass
        
        # 尝试修复markdown格式
        try:
            cleaned_text = response_text.replace('```json', '').replace('```', '').strip()
            return json.loads(cleaned_text)
        except json.JSONDecodeError:
            pass
        
        # 降级处理
        logger.warning("AI返回格式无法解析")
        return {
            "错误": "返回格式无法解析",
            "原始响应": response_text[:200] + "..." if len(response_text) > 200 else response_text,
            "笔记标题": "无法识别",
            "笔记类型": "未知",
            "原始内容": "无法识别",
            "整理后内容": "无法识别",
            "主要要点": [],
            "关键词": [],
            "重要程度": "未知",
            "建议行动": [],
            "相关主题": [],
            "笔记质量": "未知",
            "改进建议": "请检查图片质量或重新上传"
        }

    def _validate_result(self, result: Dict[str, Any]) -> bool:
        """验证结果有效性"""
        if "错误" in result:
            return False
            
        required_fields = ["笔记标题", "笔记类型", "整理后内容"]
        return all(
            field in result and 
            result[field] and 
            str(result[field]).strip() and 
            str(result[field]).strip() != "无法识别"
            for field in required_fields
        )

    def process_image(self, image_pil: Image.Image) -> Dict[str, Any]:
        """处理图像的主流程"""
        try:
            # 质量检查
            quality_info = self.check_image_quality(image_pil)
            
            # 图像预处理
            image = self.preprocess_image(image_pil)
            
            # OCR识别
            extracted_text = self.extract_text_with_ocr(image)
            
            # 显示处理信息
            if quality_info["issues"]:
                st.warning("图片质量提示：")
                for issue in quality_info["issues"]:
                    st.write(f"• {issue}")
            
            if quality_info["suggestions"]:
                st.info("改进建议：")
                for suggestion in quality_info["suggestions"]:
                    st.write(f"• {suggestion}")
            
            # AI分析
            for attempt in range(self.retry_count):
                try:
                    with st.spinner(f"正在进行第 {attempt + 1} 次分析..."):
                        result = self.analyze_note(image, extracted_text)
                    
                    if self._validate_result(result):
                        st.success(f"✅ 第 {attempt + 1} 次分析成功！")
                        return result
                    else:
                        st.warning(f"⚠️ 第 {attempt + 1} 次分析结果无效，重试中...")
                        
                except Exception as e:
                    st.error(f"❌ 第 {attempt + 1} 次尝试失败: {str(e)}")
                    if attempt < self.retry_count - 1:
                        time.sleep(2)
                        continue
            
            st.error("🚫 多次尝试后仍无法获得有效结果")
            return {"错误": "多次尝试后仍无法获得有效结果"}
            
        except Exception as e:
            logger.error(f"图像处理失败: {e}")
            return {"错误": f"图像处理失败: {str(e)}"}

    def process_multiple_images(self, images: List[Image.Image]) -> List[Dict[str, Any]]:
        """批量处理多张图片"""
        results = []
        total_images = len(images)
        
        for i, image in enumerate(images, 1):
            st.subheader(f"📄 处理第 {i}/{total_images} 页")
            result = self.process_image(image)
            results.append(result)
            
            # 显示进度
            progress = i / total_images
            st.progress(progress)
            
            # 在页面间添加分隔
            if i < total_images:
                st.markdown("---")
        
        return results

# ------------------ Streamlit 界面 ------------------
def main():
    st.set_page_config(
        page_title="📝 手写笔记识别整理系统", 
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # 侧边栏
    with st.sidebar:
        st.title("🔧 系统设置")
        st.markdown("---")
        
        # 重试次数设置
        retry_count = st.slider("重试次数", 1, 5, 3, help="设置AI分析失败时的重试次数")
        
        # 图像质量阈值
        quality_threshold = st.slider("图像质量阈值", 0, 100, 50, help="设置图像清晰度阈值")
        
        # PDF设置
        if PDF_AVAILABLE:
            pdf_scale = st.slider("PDF缩放因子", 1.0, 4.0, 2.0, 0.5, help="PDF转图像时的缩放比例，越高越清晰但处理越慢")
        
        st.markdown("---")
        st.markdown("### 📊 使用统计")
        if 'analysis_count' not in st.session_state:
            st.session_state.analysis_count = 0
        st.metric("识别次数", st.session_state.analysis_count)
        
        st.markdown("---")
        st.markdown("### 💡 使用提示")
        st.markdown("""
        1. 确保图片清晰，文字清晰可见
        2. 图片中的笔记内容应该完整显示
        3. 支持PNG、JPG、JPEG、PDF格式
        4. 如果识别失败，可以尝试重新上传
        5. 系统会自动重试，请耐心等待
        6. PDF文件会自动转换为图像进行处理
        """)
    
    # 主界面
    st.title("📝 手写笔记识别整理系统")
    st.markdown("基于文心多模态大模型，智能识别和整理手写笔记内容")
    
    # 初始化系统
    if 'system' not in st.session_state:
        st.session_state.system = NoteRecognition()
        st.session_state.system.retry_count = retry_count
    
    # 文件上传
    uploaded_file = st.file_uploader(
        "上传笔记文件", 
        type=["png", "jpg", "jpeg", "pdf"] if PDF_AVAILABLE else ["png", "jpg", "jpeg"],
        help="支持PNG、JPG、JPEG、PDF格式的文件"
    )
    
    if uploaded_file is not None:
        file_type = uploaded_file.type
        
        if file_type == "application/pdf":
            # 处理PDF文件
            if not PDF_AVAILABLE:
                st.error("PDF功能不可用，请安装PyMuPDF: pip install PyMuPDF")
            else:
                st.info("📄 检测到PDF文件，正在转换页面...")
                
                # 重置文件指针
                uploaded_file.seek(0)
                
                # 转换PDF页面为图像
                pdf_images = st.session_state.system.process_pdf(uploaded_file)
                
                if pdf_images:
                    st.success(f"✅ PDF转换成功，共 {len(pdf_images)} 页")
                    
                    # 显示PDF页面预览
                    st.subheader("📋 PDF页面预览")
                    cols = st.columns(min(3, len(pdf_images)))
                    
                    for i, img in enumerate(pdf_images):
                        col_idx = i % 3
                        with cols[col_idx]:
                            st.image(img, caption=f"第 {i+1} 页", width=300)
                    
                    # 批量处理按钮
                    if st.button("🚀 开始批量识别"):
                        st.session_state.analysis_count += 1
                        
                        with st.spinner("AI正在批量识别中，请稍候..."):
                            all_results = st.session_state.system.process_multiple_images(pdf_images)
                        
                        # 显示批量处理结果
                        st.subheader("📝 批量识别结果")
                        
                        # 创建结果汇总
                        summary = {
                            "总页数": len(all_results),
                            "成功页数": len([r for r in all_results if "错误" not in r]),
                            "失败页数": len([r for r in all_results if "错误" in r]),
                            "页面结果": all_results
                        }
                        
                        # 显示汇总信息
                        col1, col2, col3 = st.columns(3)
                        with col1:
                            st.metric("总页数", summary["总页数"])
                        with col2:
                            st.metric("成功页数", summary["成功页数"])
                        with col3:
                            st.metric("失败页数", summary["失败页数"])
                        
                        # 下载所有结果
                        summary_json = json.dumps(summary, ensure_ascii=False, indent=2)
                        st.download_button(
                            label="📥 下载所有识别结果",
                            data=summary_json,
                            file_name=f"PDF笔记识别结果_{int(time.time())}.json",
                            mime="application/json",
                            key=f"download_batch_{int(time.time())}"  # 添加唯一key
                        )
                        
                        # 显示每页详细结果
                        for i, result in enumerate(all_results):
                            with st.expander(f"第 {i+1} 页结果"):
                                if "错误" not in result:
                                    _display_single_result(result)
                                else:
                                    st.error(f"识别失败: {result['错误']}")
                else:
                    st.error("PDF转换失败，请检查文件格式")
        
        else:
            # 处理图像文件
            try:
                image = Image.open(uploaded_file)
                
                # 显示图片
                col1, col2 = st.columns([2, 1])
                with col1:
                    st.image(image, caption="上传的笔记图片")
                
                with col2:
                    st.markdown("### 📋 图片信息")
                    st.write(f"**格式**: {image.format}")
                    st.write(f"**尺寸**: {image.size[0]} × {image.size[1]}")
                    st.write(f"**模式**: {image.mode}")
                
                # 识别按钮
                if st.button("🚀 开始识别"):
                    st.session_state.analysis_count += 1
                    
                    with st.spinner("AI正在识别中，请稍候..."):
                        result = st.session_state.system.process_image(image)
                    
                    # 显示结果
                    st.subheader("📝 识别结果")
                    _display_single_result(result)
            
            except Exception as e:
                st.error(f"图片处理失败: {str(e)}")
                logger.error(f"图片处理失败: {e}")

def _display_single_result(result):
        """显示单个识别结果"""
        if "错误" not in result:
            # 成功结果
            col1, col2 = st.columns(2)
            
            with col1:
                st.metric("笔记标题", result.get('笔记标题', '未知'))
                st.metric("笔记类型", result.get('笔记类型', '未知'))
                st.metric("重要程度", result.get('重要程度', '未知'))
            
            with col2:
                st.metric("笔记质量", result.get('笔记质量', '未知'))
                st.metric("关键词数量", len(result.get('关键词', [])))
                st.metric("要点数量", len(result.get('主要要点', [])))
            
            # 详细信息
            st.markdown("### 📖 笔记内容")
            st.write(result.get('整理后内容', ''))
            
            if result.get('主要要点'):
                st.markdown("### 🎯 主要要点")
                for i, point in enumerate(result.get('主要要点', []), 1):
                    st.write(f"{i}. {point}")
            
            if result.get('关键词'):
                st.markdown("### 🔑 关键词")
                keywords = result.get('关键词', [])
                st.write(", ".join(keywords))
            
            if result.get('建议行动'):
                st.markdown("### 💡 建议行动")
                for i, action in enumerate(result.get('建议行动', []), 1):
                    st.write(f"{i}. {action}")
            
            if result.get('相关主题'):
                st.markdown("### 🏷️ 相关主题")
                topics = result.get('相关主题', [])
                st.write(", ".join(topics))
            
            if result.get('改进建议'):
                st.markdown("### 🔧 改进建议")
                st.write(result.get('改进建议', ''))
            
            # 下载结果
            result_json = json.dumps(result, ensure_ascii=False, indent=2)
            st.download_button(
                label="📥 下载识别结果",
                data=result_json,
                file_name=f"笔记识别结果_{int(time.time())}.json",
                mime="application/json",
                key=f"download_single_{hash(str(result))}"  # 添加唯一key
            )
            
        else:
            # 错误结果
            st.error(f"识别失败: {result['错误']}")
            if '原始响应' in result:
                with st.expander("查看原始响应"):
                    st.text(result['原始响应'])

if __name__ == "__main__":
    main()